"""
输入内容
时间：2024/8/26 下午5:43
"""
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from transformers import (BertTokenizer,
                          BertModel,
                          AdamW,
                          get_linear_schedule_with_warmup)
import warnings


warnings.filterwarnings('ignore')


# 自定义数据集类
class TextDataset(Dataset):
    """
    定义了一个继承自 `torch.utils.data.Dataset` 的类 `TextDataset`，用于处理文本数据。

    `__init__` 方法初始化数据集，接收文本、标签、tokenizer 和最大长度作为参数。
    `__len__` 方法返回数据集中文本的数量。
    `__getitem__` 方法对单个文本进行编码，并返回编码后的输入ID、注意力掩码和对应的标签。
    """

    def __init__(self, texts, labels, tokenizer, max_len):
        """
        初始化方法接收 texts、labels、tokenizer 和 max_len 作为参数。

        :param texts: list，文本列表。
        :param labels: list，与文本相对应的标签列表。
        :param tokenizer: `transformers` 中的 tokenizer，用于文本编码。
        :param max_len: int，每个文本的最大长度，超过该长度的文本将会被截断。
        """
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        """
        返回数据集中文本的数量，即 `self.texts` 的长度。

        :return: int，数据集中文本的数量。
        """
        return len(self.texts)

    def __getitem__(self, item):
        """
        对单个文本进行编码，并返回编码后的输入ID、注意力掩码和对应的标签。

        :param item: int，数据集中文本的索引。
        :return: dict，包含 'input_ids'、'attention_mask' 和 'labels' 键的字典。
        """
        text = str(self.texts[item])
        label = self.labels[item]

        # 使用 tokenizer 对文本进行编码
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,  # 添加特殊令牌 [CLS] 和 [SEP]
            max_length=self.max_len,  # 指定最大长度
            padding='max_length',  # 填充到最大长度
            truncation=True,  # 超过最大长度的部分将被截断
            return_attention_mask=True,  # 返回注意力掩码
            return_tensors='pt'  # 返回 PyTorch 的 tensors
        )

        # 将编码后的数据整理成模型所需的格式
        return {
            'input_ids': encoding['input_ids'].flatten(),  # 输入ID
            'attention_mask': encoding['attention_mask'].flatten(),  # 注意力掩码
            'labels': torch.tensor(label, dtype=torch.long)  # 标签
        }


# 构建分类模型
class BertClassifier(nn.Module):
    '''
    定义了一个继承自 `torch.nn.Module` 的类 `BertClassifier`，用于文本分类任务。

    `__init__` 方法初始化BERT模型和分类器。
    `forward` 方法定义了模型的前向传播过程，使用BERT模型提取文本特征，并通过分类器进行分类。
    '''
    def __init__(self, bert_model, num_classes):
        """
        初始化方法接收预训练的BERT模型和分类类别数量作为参数。

        :param bert_model: `transformers` 中的预训练BERT模型，用于提取文本特征。
        :param num_classes: int，分类任务的目标类别数量。
        """
        super(BertClassifier, self).__init__()  # 调用父类的初始化方法
        self.bert = bert_model  # BERT模型用于提取文本特征
        self.classifier = nn.Linear(768, num_classes)   # 分类器用于将特征映射到类别

    def forward(self, input_ids, attention_mask):
        """
        定义了模型的前向传播过程。

        :param input_ids: 输入文本的token ID序列，形状为 (batch_size, sequence_length)。
        :param attention_mask: 输入文本的注意力掩码，形状为 (batch_size, sequence_length)，
        用于指示哪些位置是填充的。
        :return: 输出的类别概率，形状为 (batch_size, num_classes)。
        """
        # 使用BERT模型提取文本特征
        # 注意：BERT的输出包括序列的所有token的输出和pooled_output
        # 这里我们只关心pooled_output，不关心_，它通常用于分类任务
        pooled_output = self.bert(input_ids=input_ids,
                                  attention_mask=attention_mask).pooler_output
        # print(output)
        # pooled_output = output.pooler_output
        return self.classifier(pooled_output)


def train_epoch(model, data_loader, loss_fn, optimizer, device, scheduler):
    """
        执行一个epoch的训练过程，包括前向传播、计算损失、反向传播、优化器更新等。

        :param model: torch.nn.Module，要训练的模型。
        :param data_loader: torch.utils.data.DataLoader，数据加载器，用于迭代训练数据。
        :param loss_fn: torch.nn 损失函数，用于计算预测与实际标签之间的损失。
        :param optimizer: torch.optim 优化器，用于更新模型参数。
        :param device: torch.device，指定使用的设备（CPU 或 GPU）。
        :param scheduler: 学习率调度器，用于调整优化器的学习率。
        :return: tuple，包含训练准确率和平均训练损失。
    """
    model.train()   # 将模型设置为训练模式
    losses = []     # 用于存储每个批次的损失值
    correct_predictions = 0     # 用于累计正确预测的数量

    # 遍历数据加载器中的每一个批次
    for d in data_loader:
        # 将输入数据和标签移动到指定的设备上
        input_ids = d["input_ids"].to(device)
        attention_mask = d["attention_mask"].to(device)
        labels = d["labels"].to(device)

        # 前向传播
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)

        # 计算预测的类别
        _, preds = torch.max(outputs, dim=1)

        # 计算损失
        loss = loss_fn(outputs, labels)

        # 累计正确的预测数量
        correct_predictions += torch.sum(preds == labels).item()

        # 记录损失值
        losses.append(loss.item())

        # 反向传播
        loss.backward()

        # 梯度裁剪，防止梯度爆炸
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        # 更新参数
        optimizer.step()

        # 更新学习率
        scheduler.step()

        # 清空梯度
        optimizer.zero_grad()

    # 计算准确率
    accuracy = correct_predictions / len(data_loader.dataset)

    # 计算平均损失
    avg_loss = sum(losses) / len(losses)

    return accuracy, avg_loss


def eval_model(model, data_loader, loss_fn, device):
    """
        评估模型在测试集上的性能，计算准确率和损失。

        :param model: torch.nn.Module，要评估的模型。
        :param data_loader: torch.utils.data.DataLoader，数据加载器，用于迭代测试数据。
        :param loss_fn: torch.nn 损失函数，用于计算预测与实际标签之间的损失。
        :param device: torch.device，指定使用的设备（CPU 或 GPU）。
        :return: tuple，包含测试准确率和平均测试损失。
    """
    model.eval()    # 将模型设置为评估模式
    losses = []     # 用于存储每个批次的损失值
    correct_predictions = 0     # 用于累计正确预测的数量

    # 使用 torch.no_grad() 上下文管理器禁用梯度计算
    with torch.no_grad():
        # 遍历数据加载器中的每一个批次
        for d in data_loader:
            # 将输入数据和标签移动到指定的设备上
            input_ids = d["input_ids"].to(device)
            attention_mask = d["attention_mask"].to(device)
            labels = d["labels"].to(device)

            # 前向传播
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)

            # 计算预测的类别
            _, preds = torch.max(outputs, dim=1)

            # 计算损失
            loss = loss_fn(outputs, labels)

            # 累计正确的预测数量
            correct_predictions += torch.sum(preds == labels).item()

            # 记录损失值
            losses.append(loss.item())

            # 计算准确率
        accuracy = correct_predictions / len(data_loader.dataset)

        # 计算平均损失
        avg_loss = sum(losses) / len(losses)

        return accuracy, avg_loss

# 示例数据
# texts = ["I love this movie.", "This book is terrible.",
#          "Great product!", "Service was very slow."]
# labels = [1, 0, 1, 0]  # 1: positive, 0: negative


texts = [
    "I love this movie.",
    "This book is terrible.",
    "Great product!",
    "Service was very slow.",
    "Absolutely fantastic experience!",
    "The plot was predictable and dull.",
    "Highly recommend to everyone!",
    "Terrible customer service, never again.",
    "This film is a masterpiece.",
    "The acting was subpar and unconvincing.",
    "Excellent value for the price.",
    "The delivery took too long, not happy.",
]

labels = [
    1, 0, 1, 0,  # Original data
    1, 0, 1, 0,
    1, 0, 1, 0
]


# 划分训练集和测试集
train_texts, test_texts, train_labels, test_labels \
    = train_test_split(texts, labels, test_size=0.2, random_state=42)

# 加载BERT模型和tokenizer
model_path = "G:/HUGGINGFACE_HUB_CACHE/bert-base-uncased"
tokenizer = BertTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = BertModel.from_pretrained(model_path, trust_remote_code=True)

# 初始化分类模型
num_classes = 2
classifier = BertClassifier(model, num_classes)

# 数据预处理
max_len = 64
train_dataset = TextDataset(train_texts, train_labels, tokenizer, max_len)
test_dataset = TextDataset(test_texts, test_labels, tokenizer, max_len)

# 创建DataLoader
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=2)

# 定义损失函数和优化器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
classifier.to(device)
loss_fn = nn.CrossEntropyLoss()
# optimizer = AdamW(classifier.parameters(), lr=2e-5)   # BERT模型的参数和分类器的参数一同训练
optimizer = AdamW(classifier.classifier.parameters(), lr=2e-5)

# 定义训练循环
num_epochs = 30
total_steps = len(train_loader) * num_epochs
scheduler = get_linear_schedule_with_warmup(optimizer,
                                            num_warmup_steps=0,
                                            num_training_steps=total_steps)

# 训练模型
for epoch in range(num_epochs):
    print(f'Epoch {epoch + 1}/{num_epochs}')
    print('-' * 10)

    train_acc, train_loss = train_epoch(classifier,
                                        train_loader,
                                        loss_fn,
                                        optimizer,
                                        device,
                                        scheduler)
    print(f'Train loss {train_loss} accuracy {train_acc}')

    test_acc, test_loss = eval_model(classifier, test_loader, loss_fn, device)
    print(f'Test loss {test_loss} accuracy {test_acc}')
    print()

# 保存模型状态字典
torch.save(classifier.state_dict(), '../model/model_state_dict.pth')
